Python IO - 输出重定向

今天在写 Python 单元测试的时候需要获取到待测试方法的print()输出，在Shell中使用管道重定向输出很容易，但是在Python中如何进行呢？

Python 中的标准输入输出

当我们在 Python 中打印值调用 print(value) 时候，等价于调用 sys.stdout.write(value+'\n')

print 将内容打印到了控制台，然后追加了一个换行符

input 和 sys.stdin

当我们用 raw_input(‘promption: ‘) 时，事实上是先把提示信息输出，然后捕获输入

以下两组在实际上等价：

1
2
3
4
5
hi=raw_input('Hi~')

sys.stdout.write('Hi~')
# -1 to discard the '\n' in input stream
hi=sys.stdin.readline()[:-1]

而在Python 3中新增的input()函数则在此基础上，增加了类型判断的处理，尝试将输入解析为适当的数据类型。

实际上 print 或者 raw_input 在 cpython 中的源码要复杂得多，这里‘等价’的仅是示意其与标准输入输出直接相关。源码可以参考文末给出的链接

输出重定向

默认的 sys.stdout 指向控制台，我们可以重定向到其他流(或者具有write方法的对象)：

1
2
3
4
5
6
#标准输出重定向至文件
with open('out.txt', 'w+') as f:
    sys.stdout = f 

# 标准输出重定向至字符流
sys.stdout = strout = StringIO()

完成任务后最好恢复标准输出至原先的地址(默认控制台)：

  • 可以在重定向前获取stdout的默认引用，恢复的时候再指向这个引用即可

    1
    2
    3
    4
    5
    6
    orig_stdout = sys.stdout
    # ...$重定向代码$...
    # ...print('hahaha...')...

    # 恢复
    sys.stdout = orig_stdout

  • 或使用sys.__stdout__，其保存程序开始时sys.stdout的原始值，且不受重定向影响

    1
    2
    3
    # ...$重定向代码$...
    # ...print('hahaha...')...
    sys.stdout = sys.__stdout__

结合上下文管理器，我们还可以这样：

1
2
3
4
5
6
7
8
9
10
11
12
13
@contextmanager
def stdout_redirected(new_stdout):
    save_stdout = sys.stdout
    sys.stdout = new_stdout
    try:
        yield None
    finally:
        sys.stdout = save_stdout

# 使用
with opened(filename, "w") as f:
    with stdout_redirected(f):
        print "Hello world"

应用示例：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import sys
from io import StringIO

# parse 函数直接将'解析'结果打印出来了，我们需要判断解析是否正确

def parse(query):
    print(f'query: {query}')

def test(query):
    sys.stdout = strdout = StringIO()
    parse(query)
    sys.stdout = sys.__stdout__
    return strdout.getvalue()

result = test('This message is for query')
print(result.split()[1])

# output：This

同样的，sys.stderr, sys.stdin 也都可以被重定向到多个地址。

References